"""
OpenAI compatible API data models for ChatCompletion endpoints
"""
from typing import Optional, List, Dict, Any, Union, Literal
from pydantic import BaseModel, Field
from enum import Enum


class MessageRole(str, Enum):
    """Message roles for chat completion"""
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    FUNCTION = "function"
    TOOL = "tool"


class FunctionCall(BaseModel):
    """Function call in message"""
    name: str = Field(..., description="Function name")
    arguments: str = Field(..., description="Function arguments as JSON string")


class ToolCall(BaseModel):
    """Tool call in message"""
    id: str = Field(..., description="Tool call ID")
    type: Literal["function"] = Field("function", description="Tool call type")
    function: FunctionCall = Field(..., description="Function call details")


class Message(BaseModel):
    """Chat message format"""
    role: MessageRole = Field(..., description="Message role")
    content: Optional[str] = Field(None, description="Message content")
    name: Optional[str] = Field(None, description="Message name (for function/tool messages)")
    function_call: Optional[FunctionCall] = Field(None, description="Function call (deprecated)")
    tool_calls: Optional[List[ToolCall]] = Field(None, description="Tool calls")
    tool_call_id: Optional[str] = Field(None, description="Tool call ID (for tool role)")

    class Config:
        use_enum_values = True


class Function(BaseModel):
    """Function definition for function calling"""
    name: str = Field(..., description="Function name")
    description: Optional[str] = Field(None, description="Function description")
    parameters: Dict[str, Any] = Field(..., description="Function parameters schema")


class Tool(BaseModel):
    """Tool definition"""
    type: Literal["function"] = Field("function", description="Tool type")
    function: Function = Field(..., description="Function definition")


class ResponseFormat(BaseModel):
    """Response format specification"""
    type: Literal["text", "json_object"] = Field("text", description="Response format type")


class StreamOptions(BaseModel):
    """Stream options"""
    include_usage: Optional[bool] = Field(False, description="Include usage in stream")


class OpenAIRequest(BaseModel):
    """OpenAI ChatCompletion request format"""
    model: str = Field(..., description="Model identifier")
    messages: List[Message] = Field(..., description="Chat messages")
    temperature: Optional[float] = Field(1.0, ge=0.0, le=2.0, description="Sampling temperature")
    top_p: Optional[float] = Field(1.0, ge=0.0, le=1.0, description="Nucleus sampling")
    n: Optional[int] = Field(1, ge=1, description="Number of completions")
    stream: Optional[bool] = Field(False, description="Stream partial results")
    stop: Optional[Union[str, List[str]]] = Field(None, description="Stop sequences")
    max_tokens: Optional[int] = Field(None, ge=1, description="Maximum tokens to generate")
    presence_penalty: Optional[float] = Field(0.0, ge=-2.0, le=2.0, description="Presence penalty")
    frequency_penalty: Optional[float] = Field(0.0, ge=-2.0, le=2.0, description="Frequency penalty")
    logit_bias: Optional[Dict[str, float]] = Field(None, description="Logit bias")
    user: Optional[str] = Field(None, description="User identifier")
    
    # Function calling (deprecated)
    functions: Optional[List[Function]] = Field(None, description="Available functions (deprecated)")
    function_call: Optional[Union[str, Dict[str, str]]] = Field(None, description="Function call control (deprecated)")
    
    # Tool calling (preferred)
    tools: Optional[List[Tool]] = Field(None, description="Available tools")
    tool_choice: Optional[Union[str, Dict[str, Any]]] = Field(None, description="Tool choice control")
    
    # Additional parameters
    response_format: Optional[ResponseFormat] = Field(None, description="Response format")
    seed: Optional[int] = Field(None, description="Random seed for reproducible outputs")
    stream_options: Optional[StreamOptions] = Field(None, description="Stream options")


class Usage(BaseModel):
    """Token usage statistics"""
    prompt_tokens: int = Field(..., description="Tokens in prompt")
    completion_tokens: int = Field(..., description="Tokens in completion")
    total_tokens: int = Field(..., description="Total tokens used")


class Choice(BaseModel):
    """Completion choice"""
    index: int = Field(..., description="Choice index")
    message: Message = Field(..., description="Completion message")
    finish_reason: Optional[Literal["stop", "length", "function_call", "tool_calls", "content_filter"]] = Field(
        None, description="Reason for completion finish"
    )


class OpenAIResponse(BaseModel):
    """OpenAI ChatCompletion response format"""
    id: str = Field(..., description="Completion ID")
    object: Literal["chat.completion"] = Field("chat.completion", description="Object type")
    created: int = Field(..., description="Creation timestamp")
    model: str = Field(..., description="Model used")
    choices: List[Choice] = Field(..., description="Completion choices")
    usage: Usage = Field(..., description="Token usage")
    system_fingerprint: Optional[str] = Field(None, description="System fingerprint")


class StreamChoice(BaseModel):
    """Stream completion choice"""
    index: int = Field(..., description="Choice index")
    delta: Message = Field(..., description="Delta message")
    finish_reason: Optional[Literal["stop", "length", "function_call", "tool_calls", "content_filter"]] = Field(
        None, description="Reason for completion finish"
    )


class StreamResponse(BaseModel):
    """OpenAI ChatCompletion stream response format"""
    id: str = Field(..., description="Completion ID")
    object: Literal["chat.completion.chunk"] = Field("chat.completion.chunk", description="Object type")
    created: int = Field(..., description="Creation timestamp")
    model: str = Field(..., description="Model used")
    choices: List[StreamChoice] = Field(..., description="Stream choices")
    usage: Optional[Usage] = Field(None, description="Token usage (final chunk only)")


class OpenAIError(BaseModel):
    """OpenAI API error format"""
    message: str = Field(..., description="Error message")
    type: str = Field(..., description="Error type")
    param: Optional[str] = Field(None, description="Parameter that caused error")
    code: Optional[str] = Field(None, description="Error code")


class OpenAIErrorResponse(BaseModel):
    """OpenAI API error response format"""
    error: OpenAIError = Field(..., description="Error details")


class ModelInfo(BaseModel):
    """Model information"""
    id: str = Field(..., description="Model ID")
    object: Literal["model"] = Field("model", description="Object type")
    created: int = Field(..., description="Creation timestamp")
    owned_by: str = Field(..., description="Model owner")


class ModelsResponse(BaseModel):
    """Models list response"""
    object: Literal["list"] = Field("list", description="Object type")
    data: List[ModelInfo] = Field(..., description="Available models")